[Memory optm] loss using torch + compile #337
Conversation
meta-cla
bot
added
the
CLA Signed
| return logprobs | ||
|
|
||
| # Convert to fp32 for numerical stability | ||
| scaled_logits_fp32 = scaled_logits.float() |
There was a problem hiding this comment.
Noob question: what's the dtype for scaled_logits?
There was a problem hiding this comment.
float becomes torch.float32
|
@ebsmothers @Jack-Khuu @joecummings @pbontrager can some of you confirm that i dont need to do the all_gather that was happening in |
| import torch.nn.functional as F | ||
|
|
||
|
|
||
| def selective_log_softmax(logits: torch.Tensor, index: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
Should we also delete this function?
There was a problem hiding this comment.
its used in 3 other places. I will leave it there for now. We will prob need some larger refactoring later to clean up / organize losses
|
|
Co-authored-by: Jiyue Wang <[email protected]>
|
|
||
| # compile loss | ||
| logger.info("Compiling loss") | ||
| self.loss = torch.compile(self.loss) |
There was a problem hiding this comment.
Is there any circumstance under which this command would fail?
There was a problem hiding this comment.
cant think of one in our scenario, but if/when this happens, we can fix it
| logprobs = selective_log_softmax(scaled_logits, input_ids) | ||
| return logprobs | ||
|
|
||
| # Cast up to fp32 for numerical stability |
There was a problem hiding this comment.
nit: I would change this to something like "ensure logits are in fp32" b/c they actually could already be in fp32 and no need for "Casting up"
Co-authored-by: Joe Cummings <[email protected]>
…into compile_loss
4 participants
Memory freebies
I dont think that loss/reward is a good way to check correctness here. But i compared the functions locally and they provide the same output.